import json
import asyncio
import logging
from datetime import datetime
from typing import List, Dict, Any

import redis
from pymilvus import MilvusClient

from core.config import settings
from core.llm import get_default_llm, get_embedding_llm

logger = logging.getLogger(__name__)

class ChatService:
    def __init__(self):
        self.milvus_client = MilvusClient(settings.MILVUS_URL)
        
        # 尝试连接Redis，如果失败则使用内存缓存
        try:
            self.redis_client = redis.Redis(
                host=settings.REDIS_HOST, 
                port=settings.REDIS_PORT, 
                db=settings.REDIS_DB, 
                password=settings.REDIS_PASSWORD,
                decode_responses=True,  # 自动解码
                socket_connect_timeout=5,  # 连接超时
                socket_timeout=5 # 读写超时
            )
            # 测试连接
            self.redis_client.ping()
            self.redis_available = True
            logger.info("Redis连接成功")
        except Exception as e:
            logger.warning(f"Redis连接失败，使用内存缓存: {e}")
            self.redis_client = None
            self.redis_available = False
        
        self.llm = get_default_llm()
        self.embedding_llm = get_embedding_llm()
        
        # 添加缓存
        self._embedding_cache = {}
        self._summary_cache = {}
        
    def _get_chat_key(self, account_id: str) -> str:
        return f"chat_history:{account_id}"
    
    def save_chat_history(self, account_id: str, message: List[Dict[str, str]]) -> None:
        if not self.redis_available:
            # 使用内存缓存
            key = self._get_chat_key(account_id)
            if not hasattr(self, '_chat_history_cache'):
                self._chat_history_cache = {}
            if key not in self._chat_history_cache:
                self._chat_history_cache[key] = []
            self._chat_history_cache[key].insert(0, message)
            return
            
        key = self._get_chat_key(account_id)
        self.redis_client.lpush(key, json.dumps(message))
        
    def get_chat_history(self, account_id: str) -> List[Dict[str, str]]:
        if not self.redis_available:
            # 使用内存缓存
            key = self._get_chat_key(account_id)
            if hasattr(self, '_chat_history_cache') and key in self._chat_history_cache:
                return self._chat_history_cache[key]
            return []
            
        key = self._get_chat_key(account_id)
        messages = self.redis_client.lrange(key, 0, -1)
        if messages:
            return [json.loads(message) for message in messages]
        return []
    
    def add_message(self, account_id: str, role: str, content: str) -> None:
        message = self.get_chat_history(account_id)
        message.append({"role": role, "content": content, "timestamp": datetime.now().isoformat()})
        self.save_chat_history(account_id, message)
      
    def save_chat_message(self, account_id: str, user_message: str, assistant_message: str):
        self.add_message(account_id, "user", user_message)
        self.add_message(account_id, "assistant", assistant_message)
      
    # 异步保存聊天消息，不阻塞响应
    async def save_chat_message_async(self, account_id: str, user_message: str, assistant_message: str):
        await asyncio.get_event_loop().run_in_executor(
            None, self.save_chat_message, account_id, user_message, assistant_message
        )
      
    def clear_chat_history(self, account_id: str) -> None:
        if self.redis_available:
            key = self._get_chat_key(account_id)
            self.redis_client.delete(key)
        else:
            # 清除内存缓存
            key = self._get_chat_key(account_id)
            if hasattr(self, '_chat_history_cache') and key in self._chat_history_cache:
                del self._chat_history_cache[key]
        
    def get_summary_key(self, account_id: str) -> str:
        return f"summary:{account_id}"

    def get_information(self, message: str, course_id: str):
        # 检查缓存
        cache_key = f"{message}:{course_id}"
        if cache_key in self._embedding_cache:
            return self._embedding_cache[cache_key]
            
        try:
            # TODO 后续用多查询检索增强进行优化
            embedding = self.embedding_llm.embed_query(message)
            result = self.milvus_client.search(
                collection_name="teach_platform",
                data=[embedding],
                output_fields=["content"],
                filter=f"course_id=='{course_id}'",
                limit=5,  # 限制返回结果数量
            )
            
            # 缓存结果
            self._embedding_cache[cache_key] = result
            return result
        except Exception as e:
            logger.warning(f"向量检索失败，返回空结果: {e}")
            # 返回空结果，避免阻塞
            return []
    
    # 异步获取信息
    async def get_information_async(self, message: str, course_id: str):
        return await asyncio.get_event_loop().run_in_executor(
            None, self.get_information, message, course_id
        )
        
    async def generate_summary(self, account_id: str) -> str:
        try:
            # 检查缓存
            if account_id in self._summary_cache:
                return self._summary_cache[account_id]
                
            messages = self.get_chat_history(account_id)
            if not messages:
                return ""
                
            # 限制消息数量，避免过长
            if len(messages) > 20:
                messages = messages[-20:]
                
            prompt = f"""请根据以下对话生成一个简要的核心摘要，突出话题和关键信息：
            
            {json.dumps(messages, ensure_ascii=False, indent=2)}
            
            摘要要求：
            1.突出对话的主要内容和关键信息
            2.使用第三人称描述，提取重要数据/时间节点/待办事项
            3.保留原始对话中的重要细节
            4.确保包含最新的对话内容
            
            """ 
            
            response = await self.llm.ainvoke(prompt)
            new_summary = response.content
            
            try:
                if self.redis_available:
                    summary_key = self.get_summary_key(account_id)
                    old_summary = self.redis_client.get(summary_key)
                    
                    if old_summary:
                        merged_prompt = f"""请将以下两个摘要合并成新的摘要：
                        旧摘要：{old_summary}
                        新摘要：{new_summary}
                        
                        合并要求：
                        1.保留两个摘要中的重要信息
                        2.突出对话的主要话题和关键信息
                        3.使用第三人称描述，提取重要数据/时间节点/待办事项
                        4.保留原始对话中的重要细节
                        5.确保包含最新的对话内容
                        """
                        merged_response = await self.llm.ainvoke(merged_prompt)
                        final_summary = merged_response.content
                    else:
                        final_summary = new_summary
                    
                    # 异步保存到Redis
                    await asyncio.get_event_loop().run_in_executor(
                        None, self.redis_client.set, summary_key, final_summary
                    )
                else:
                    # 使用内存缓存
                    summary_key = self.get_summary_key(account_id)
                    if not hasattr(self, '_summary_cache_dict'):
                        self._summary_cache_dict = {}
                    
                    old_summary = self._summary_cache_dict.get(summary_key)
                    if old_summary:
                        merged_prompt = f"""请将以下两个摘要合并成新的摘要：
                        旧摘要：{old_summary}
                        新摘要：{new_summary}
                        
                        合并要求：
                        1.保留两个摘要中的重要信息
                        2.突出对话的主要话题和关键信息
                        3.使用第三人称描述，提取重要数据/时间节点/待办事项
                        4.保留原始对话中的重要细节
                        5.确保包含最新的对话内容
                        """
                        merged_response = await self.llm.ainvoke(merged_prompt)
                        final_summary = merged_response.content
                    else:
                        final_summary = new_summary
                    
                    self._summary_cache_dict[summary_key] = final_summary
            except Exception as e:
                logger.warning(f"摘要操作失败，使用新摘要: {e}")
                final_summary = new_summary
            
            # 缓存摘要
            self._summary_cache[account_id] = final_summary
            
            return final_summary
            
        except Exception as e:
            logger.error(f"生成摘要失败: {e}")
            return ""
